# -*- coding: utf-8 -*-

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split

# 假设我们有一个包含多个股票价格数据的DataFrame，每列代表一个特征变量
# 示例数据生成
np.random.seed(42)
dates = pd.date_range('2020-01-01', '2021-01-01', freq='B')  # 生成工作日日期
num_stocks = 5
num_features = 5  # 开盘价、收盘价、最高价、最低价、交易量
data = np.random.rand(len(dates), num_stocks * num_features) * 100  # 随机生成特征数据
columns = [f'Stock_{i+1}_{feature}' for i in range(num_stocks) for feature in ['Open', 'Close', 'High', 'Low', 'Volume']]
df = pd.DataFrame(data, columns=columns, index=dates)

# 数据预处理
scaler = StandardScaler()
scaled_data = scaler.fit_transform(df.values)

# 将数据转换为适合Transformer输入的格式
sequence_length = 30  # 使用过去30天的数据预测第31天的价格
X = []
y = []
for i in range(len(scaled_data) - sequence_length):
    X.append(scaled_data[i:i+sequence_length])
    # 选取每只股票的收盘价：每只股票占5个特征，索引为1, 6, 11, ... (即 1 + 5*j)
    target = scaled_data[i+sequence_length][1::5]  # (num_stocks,)
    y.append(target)
X = np.array(X)
y = np.array(y)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)

# 转换为PyTorch张量
X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
y_train_tensor = torch.tensor(y_train, dtype=torch.float32)
X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
y_test_tensor = torch.tensor(y_test, dtype=torch.float32)

# Transformer模型定义
class TransformerModel(nn.Module):
    def __init__(self, input_dim, model_dim, num_heads, num_layers, output_dim):
        super(TransformerModel, self).__init__()
        self.model_dim = model_dim
        self.input_fc = nn.Linear(input_dim, model_dim)
        self.transformer = nn.Transformer(d_model=model_dim, nhead=num_heads, num_encoder_layers=num_layers)
        self.output_fc = nn.Linear(model_dim, output_dim)
        self.pos_encoder = nn.Embedding(1000, model_dim)  # 假设最大序列长度为1000

    def forward(self, x):
        seq_len, batch_size, _ = x.size()
        positions = torch.arange(0, seq_len, dtype=torch.long).unsqueeze(1).expand(seq_len, batch_size)
        x = self.input_fc(x) + self.pos_encoder(positions)
        x = x.permute(1, 0, 2)  # 转换为 (batch_size, seq_len, model_dim)
        x = self.transformer(x, x)
        x = x.mean(dim=1)  # 对序列维度取平均
        x = self.output_fc(x)
        return x
    
# 模型实例化
input_dim = num_stocks * num_features  # 每个时间步的特征数量
model_dim = 64
num_heads = 4
num_layers = 2
output_dim = num_stocks  # 预测每个股票的收盘价（即每只股票的一个值）
model = TransformerModel(input_dim, model_dim, num_heads, num_layers, output_dim)

# 训练设置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练过程
num_epochs = 50
for epoch in range(num_epochs):
    model.train()
    optimizer.zero_grad()
    output = model(X_train_tensor.to(device))
    loss = criterion(output, y_train_tensor.to(device))
    loss.backward()
    optimizer.step()
    if (epoch+1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

# 测试过程
model.eval()
with torch.no_grad():
    y_pred = model(X_test_tensor.to(device))
    y_pred = y_pred.cpu().numpy()
    y_test = y_test_tensor.numpy()

# 评估模型
mse = mean_squared_error(y_test, y_pred)
print(f'Mean Squared Error on Test Set: {mse:.4f}')